import math
import cv2 as cv
import numpy as np
from scipy import signal

'''
函数描述：原始代码，增加了二值化。
JZhou@20211125
'''

def pascal_smooth(n):
  	# 返回 n 阶的非归一化的高斯平滑算子
    pascal_smooth = np.zeros([1, n], np.float32)
    for i in range(n):
        pascal_smooth[0][i] = math.factorial(n-1) / (math.factorial(i) * math.factorial(n-1-i))
    return pascal_smooth

def pascal_diff(n):
  	# 返回 n 阶的差分算子
    pascal_diff = np.zeros([1, n], np.float32)
    pascal_smooth_previous = pascal_smooth(n-1)
    for i in range(n):
        if i == 0:
            # 恒等于 1
            pascal_diff[0][i] = pascal_smooth_previous[0][i]
        elif i == n-1:
            # 恒等于 -1
            pascal_diff[0][i] = - pascal_smooth_previous[0][i-1]
        else:
            pascal_diff[0][i] = pascal_smooth_previous[0][i] - pascal_smooth_previous[0][i-1]
    return pascal_diff


def get_sobel_kernel(n):
    pascal_smooth_kernel = pascal_smooth(n)
    pascal_diff_kernel = pascal_diff(n)
    # 水平方向的卷积核
    sobel_kerenl_x = signal.convolve2d(pascal_smooth_kernel.transpose(), pascal_diff_kernel, mode='full')
    # 垂直方向的卷积核
    sobel_kerenl_y = signal.convolve2d(pascal_smooth_kernel, pascal_diff_kernel.transpose(), mode='full')
    return sobel_kerenl_x, sobel_kerenl_y

def sobel(img, n):
    rows, cols = img.shape[:2]
    # 平滑算子
    pascal_smooth_kernel = pascal_smooth(n)
    # 差分算子
    pascal_diff_kernel = pascal_diff(n)
    # 水平方向上的 sobel 核卷积
    # 先进行垂直方向的平滑
    img_sobel_x = signal.convolve2d(img, pascal_smooth_kernel.transpose(), mode='same')
    # 再进行水平方向上的差分
    img_sobel_x = signal.convolve2d(img_sobel_x, pascal_diff_kernel, mode='same')
    # 垂直方向上的 sobel 核卷积
    img_sobel_y = signal.convolve2d(img, pascal_smooth_kernel, mode='same')
    img_sobel_y = signal.convolve2d(img_sobel_y, pascal_diff_kernel.transpose(), mode='same')

    return img_sobel_x, img_sobel_y


def handlePic(img, save_path):
    blurred = cv.pyrMeanShiftFiltering(img, 30, 30) # 均值迁移去噪声
    img_gray = cv.cvtColor(blurred,cv.COLOR_RGB2GRAY) # 将图片转为灰度图
    img_sobel_x, img_sobel_y = sobel(img_gray, 3)

    img_sobel_x_c, img_sobel_y_c = img_sobel_x.copy(), img_sobel_y.copy()
    img_sobel_x_c, img_sobel_y_c = abs(img_sobel_x_c), abs(img_sobel_y_c)
    img_sobel_x_c[img_sobel_x_c>255] = 255
    img_sobel_y_c[img_sobel_y_c>255] = 255
    img_sobel_x_c = img_sobel_x_c.astype(np.uint8)
    img_sobel_y_c = img_sobel_y_c.astype(np.uint8)
    # 平方和开方的方式
    edge = np.sqrt(np.power(img_sobel_x, 2.0) + np.power(img_sobel_y, 2.0))

    # 直接截断显示
    edge_c = edge.copy()
    edge_c[edge_c > 255] = 255
    edge_c = edge_c.astype(np.uint8)
    cv.imwrite('/Users/vine/Desktop/edge_c.jpg', edge_c)
    ret2,th2 = cv.threshold(edge_c,245,255,cv.THRESH_BINARY)
    cv.imwrite('/Users/vine/Desktop/th2.jpg', th2)


if __name__ == '__main__':
    # imgInit = cv.imread('/Users/vine/Documents/Winglab/ISPTestCode/testPic/Clip_V1/2021091120354001191G-CL0615-6.jpg')

    edge_c = cv.imread("/Users/vine/Desktop/edge_c.jpg") # 读入钢片边缘
    th2 = cv.imread("/Users/vine/Desktop/th2.jpg") # 读入整体边缘

    # 提取边缘图中的点集
    points = []
    img_info = edge_c.shape
    rows = img_info[0]
    cols = img_info[1]
    for i in range(rows):
        for j in range(cols):
            k = edge_c[i, j] # 二值图的像素点
            # print(k)
            if int(k[0]) + int(k[1]) + int(k[2]) > 600:
                points.append([i, j])

    # 把边缘图和原图重合
    img_array = np.array(th2)
    img2_array = img_array
    for p in points:
        img2_array[p[0],p[1]]=np.array([0, 255, 255])

    cv.imwrite("/Users/vine/Desktop/img2_array.jpg", img2_array)
    

